import os
import dlib
import cv2


class dispose():
    def __init__(self):
        self.detection_model = dlib.get_frontal_face_detector()

    def dispose(self, in_path, out_path):

        img = cv2.imread(in_path)
        dets, confidence, direction = self.detection_model.run(img)
        if len(dets) > 0:
            # for i, d in enumerate(dets):
            #     x1 = d.left()
            #     y1 = d.top()
            #     x2 = d.right()
            #     y2 = d.bottom()
            #     if x1 < 0:
            #         x1 = 0
            #     if y1 < 0:
            #         y1 = 0
            #     out_path2 = os.path.splitext(out_path)
            #     out_path3 = out_path2[0] + '_' + str(i) + out_path2[1]
            #     cv2.imwrite(out_path3, img[y1:y2, x1:x2])
            #     print(out_path3)
            d = dets[0]
            x1 = d.left()
            y1 = d.top()
            x2 = d.right()
            y2 = d.bottom()
            if x1 < 0:
                x1 = 0
            if y1 < 0:
                y1 = 0
            cv2.imwrite(out_path, img[y1:y2, x1:x2])

    # 遍历文件
    def traverse_dir(self, in_path, out_path):
        for path in os.listdir(in_path):
            in_sub_path = os.path.join(in_path, path)
            out_sub_path = os.path.join(out_path, path)
            if os.path.isdir(in_sub_path):
                if not os.path.isdir(out_sub_path):
                    os.makedirs(out_sub_path)
                self.traverse_dir(in_sub_path, out_sub_path)
            else:
                self.dispose(in_sub_path, out_sub_path)


if __name__ == '__main__':

    dis = dispose()

    input_root_path = '/home/linux/database/flw/lfw'
    output_root_path = '/home/linux/database/flw/output'

    if not os.path.isdir(input_root_path):
        print('输入目录未找到')
        exit(0)

    if not os.path.isdir(output_root_path):
        os.makedirs(output_root_path)

    dis.traverse_dir(input_root_path, output_root_path)
